package com.hzsparrow.ai.controller;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import com.hzsparrow.ai.config.PromptConfiguration;
import com.hzsparrow.ai.entity.ChatHistory;
import com.hzsparrow.ai.entity.Conversation;
import com.hzsparrow.ai.entity.KnowledgeBase;
import com.hzsparrow.ai.entity.User;
import com.hzsparrow.ai.service.ChatHistoryService;
import com.hzsparrow.ai.service.ChatService;
import com.hzsparrow.ai.service.ConversationService;
import com.hzsparrow.ai.service.DocumentService;
import com.hzsparrow.ai.service.KnowledgeBaseService;
import com.hzsparrow.ai.service.UserService;

import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContextHolder;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;

@RestController
public class ChatController {

    private final ChatService chatService;
    private final DocumentService documentService;
    private final ChatHistoryService chatHistoryService;
    private final KnowledgeBaseService knowledgeBaseService;
    private final UserService userService;
    private final ConversationService conversationService;

    @Autowired
    private VectorStore vectorStore;

    @Autowired
    private OllamaChatModel chatModel;

    @Autowired
    private PromptConfiguration promptConfig;

    @Autowired
    public ChatController(
            ChatService chatService, 
            DocumentService documentService,
            ChatHistoryService chatHistoryService,
            KnowledgeBaseService knowledgeBaseService,
            UserService userService,
            ConversationService conversationService) {
        this.chatService = chatService;
        this.documentService = documentService;
        this.chatHistoryService = chatHistoryService;
        this.knowledgeBaseService = knowledgeBaseService;
        this.userService = userService;
        this.conversationService = conversationService;
    }

    @GetMapping("/api/chat1")
    public Flux<String> chat1(@RequestParam String message) {
        // 首先尝试从向量库检索相关文档
        List<Document> relevantDocs = vectorStore.similaritySearch(
                SearchRequest.builder()
                        .query(message)
                        .similarityThreshold(0.6d)
                        .topK(6)
                        .build());

        // 创建ChatClient构建器
        ChatClient.Builder chatClientBuilder = ChatClient.builder(chatModel);

        // 判断是否找到相关文档
        if (!relevantDocs.isEmpty()) {
            // 找到相关文档，使用RAG方式回答
            System.out.println("找到相关文档，使用RAG方式回答，文档数量: " + relevantDocs.size());

            // 创建QuestionAnswerAdvisor用于RAG
            QuestionAnswerAdvisor questionAnswerAdvisor = QuestionAnswerAdvisor.builder(vectorStore)
                    .searchRequest(SearchRequest.builder()
                            .similarityThreshold(0.6d)
                            .topK(6)
                            .build())
                    .userTextAdvise(promptConfig.getVectorChatPrompt())
                    .build();

            // 返回RAG方式的回答
            return chatClientBuilder.build().prompt()
                    .advisors(questionAnswerAdvisor)
                    .user(message)
                    .stream()
                    .content();
        } else {
            // 未找到相关文档，使用普通方式回答
            System.out.println("未找到相关文档，使用普通方式回答");

            // 创建包含系统提示和用户提示的Prompt
            Prompt prompt = new Prompt(
                    new SystemMessage(promptConfig.getGeneralChatPrompt()),
                    new UserMessage(message)
            );

            // 返回普通方式的回答
            return chatClientBuilder.build().prompt(prompt)
                    .stream()
                    .content();
        }

    }

    @GetMapping("/api/chat")
    public Flux<String> chat(
            @RequestParam String message,
            @RequestParam(required = false) Long knowledgeBaseId,
            @RequestParam(required = false) String sessionId) {
        
        // 获取当前用户（如果已认证）
        User currentUser = null;
        try {
            Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
            if (authentication != null && authentication.isAuthenticated() && 
                    !authentication.getName().equals("anonymousUser")) {
                currentUser = userService.getUserByUsername(authentication.getName());
            }
        } catch (Exception e) {
            // 未认证用户，忽略错误继续处理
        }
        
        // 保存变量的最终引用，以便在lambda中使用
        final User finalCurrentUser = currentUser;
        
        // 创建或获取会话ID
        final String finalSessionId = (sessionId == null || sessionId.isEmpty()) 
                ? UUID.randomUUID().toString() 
                : sessionId;
        
        // 检查知识库是否存在
        KnowledgeBase knowledgeBase = null;
        if (knowledgeBaseId != null) {
            knowledgeBase = knowledgeBaseService.getKnowledgeBase(knowledgeBaseId);
        }
        
        // 保存变量的最终引用，以便在lambda中使用
        final KnowledgeBase finalKnowledgeBase = knowledgeBase;
        final Long finalKnowledgeBaseId = knowledgeBaseId;
        
        // 获取向量索引名称，如果指定了知识库
        String vectorIndexName = finalKnowledgeBase != null ? finalKnowledgeBase.getVectorIndexName() : null;
        
        // 检查或创建Conversation对象
        if (finalCurrentUser != null) {
            // 检查是否存在此会话ID的对话
            Conversation existingConversation = conversationService.getConversationBySessionId(finalSessionId);
            
            if (existingConversation == null) {
                // 创建新的对话记录
                Conversation newConversation = new Conversation();
                newConversation.setSessionId(finalSessionId);
                newConversation.setUserId(finalCurrentUser.getId());
                newConversation.setKnowledgeBaseId(finalKnowledgeBaseId);
                
                // 使用消息内容的前30个字符作为对话标题
                String title = message.length() > 30 ? message.substring(0, 30) : message;
                newConversation.setTitle(title);
                
                conversationService.createConversation(newConversation);
            } else {
                // 更新对话的最后修改时间
                conversationService.updateConversation(existingConversation);
            }
        }
        
        // 处理聊天请求，同时记录对话历史
        Flux<String> response = chatService.processChat(message, vectorIndexName);
        
        // 转换为字符串以保存对话历史
        StringBuilder responseBuilder = new StringBuilder();
        
        return response
                .doOnNext(chunk -> {
                    // 收集响应内容
                    responseBuilder.append(chunk);
                })
                .doOnComplete(() -> {
                    String responseText = responseBuilder.toString();
                    saveHistoryAsync(finalCurrentUser, finalSessionId, finalKnowledgeBaseId, 
                    message, responseText, finalKnowledgeBase != null);
                });
    }

    private void saveHistoryAsync(User user, String sessionId, Long knowledgeBaseId, 
                           String message, String response, boolean usedKnowledgeBase) {
    // 使用异步但非反应式的方式保存历史
    CompletableFuture.runAsync(() -> {
        try {
            if (user != null) {
                ChatHistory history = new ChatHistory();
                history.setSessionId(sessionId);
                history.setUserId(user.getId());
                history.setKnowledgeBaseId(knowledgeBaseId);
                history.setUserMessage(message);
                history.setAiResponse(response);
                history.setUsedKnowledgeBase(usedKnowledgeBase);
                chatHistoryService.save(history); // 使用基础的MyBatis-Plus保存方法
            }
        } catch (Exception e) {
            System.err.println("异步保存聊天历史失败: " + e.getMessage());
        }
    });
}

    // 保留原来的GET方法以保持向后兼容
    @GetMapping("/api/addDocument")
    public String addDocument(@RequestParam String content) {
        return documentService.addDocument(content, null, "未命名文档");
    }
    
    // POST方法用于添加文档
    @PostMapping("/api/addDocument")
    public String addDocumentPost(@RequestBody DocumentRequest request) {
        return documentService.addDocument(request.getContent(), null, "未命名文档");
    }
    
    // 内部类，用于接收POST请求的JSON数据
    public static class DocumentRequest {
        private String content;
        
        public String getContent() {
            return content;
        }
        
        public void setContent(String content) {
            this.content = content;
        }
    }
}